package io.sunny.platform.ai.deepseek.service;

import cn.hutool.core.thread.NamedThreadFactory;
import io.github.pigmesh.ai.deepseek.core.DeepSeekClient;
import io.github.pigmesh.ai.deepseek.core.chat.AssistantMessage;
import io.github.pigmesh.ai.deepseek.core.chat.ChatCompletionChoice;
import io.github.pigmesh.ai.deepseek.core.chat.ChatCompletionRequest;
import io.github.pigmesh.ai.deepseek.core.chat.ChatCompletionResponse;
import io.sunny.platform.ai.deepseek.config.DeepSeekConfigProperties;
import io.sunny.platform.ai.deepseek.processor.BatchProcessor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cache.Cache;
import org.springframework.cache.CacheManager;
import org.springframework.cache.annotation.Cacheable;
import org.springframework.stereotype.Service;
import org.springframework.util.DigestUtils;

import java.time.Duration;
import java.util.List;
import java.util.concurrent.*;

/**
 * @Author SUNNY
 * @Date 2025/11/30 11:10
 * @PackageName: io.sunny.platform.ai.deepseek.service
 * @Description: 智能批处理与缓存策略
 * @Version 1.0
 */
@Slf4j
@Service
public class AIPerformanceOptimizer {
    @Autowired
    private CacheManager cacheManager;
    private final DeepSeekClient deepSeekClient;
    private final BatchProcessor<ChatCompletionRequest, ChatCompletionResponse> batchProcessor;

    public AIPerformanceOptimizer(DeepSeekConfigProperties config) {
        this.batchProcessor = new BatchProcessor.Builder<ChatCompletionRequest, ChatCompletionResponse>()
                .batchSize(10)                       // 每批10个请求
                .bufferSize(100)                     // 缓冲区大小
                .maxWaitTime(Duration.ofMillis(50))  // 最长等待时间
                .processor(this::processBatch)
                .build();
        this.deepSeekClient = new DeepSeekClient.Builder()
                .openAiApiKey(config.getKey())
                .baseUrl(config.getBaseUrl())
                .connectTimeout(config.getTimeout().getConnect())
                .readTimeout(config.getTimeout().getRead())
                .build();
    }

    @Cacheable(value = "aiResponses",
            key = "T(com.example.util.CacheKeyGenerator).generate(#request)",
            unless = "#result == null || #result.fallback")
    public ChatCompletionResponse getOptimizedResponse(ChatCompletionRequest request) {
        // 先检查本地缓存
        String cacheKey = generateCacheKey(request);
        Cache cache = cacheManager.getCache("aiResponses");
        Cache.ValueWrapper cached = cache.get(cacheKey);

        if (cached != null) {
            log.debug("从缓存中获取结果：{}", cacheKey);
            return (ChatCompletionResponse) cached.get();
        }

        // 批处理执行
        batchProcessor.submit(request);

        for (;;) {
            cached = cache.get(cacheKey);
            if (cached != null) {
                log.debug("从缓存中获取结果：{}", cacheKey);
                return (ChatCompletionResponse) cached.get();
            }
        }
    }

    private List<ChatCompletionResponse> processBatch(List<ChatCompletionRequest> requests) {
        long startTime = System.currentTimeMillis();

        try {
            // 并行处理请求
            List<CompletableFuture<ChatCompletionResponse>> futures = requests.stream()
                    .map(req -> CompletableFuture.supplyAsync(() ->
                            deepSeekClient.chatCompletion(req).execute(),
                            getAIPoolExecutor()))
                    .toList();

            List<ChatCompletionResponse> responses = futures.stream()
                    .map(CompletableFuture::join)
                    .toList();

            long costTime = System.currentTimeMillis() - startTime;
            log.info("批处理完成，数量：{}，耗时：{}ms", requests.size(), costTime);

            // 缓存结果
            cacheResponses(requests, responses);

            return responses;
        } catch (Exception e) {
            log.error("批处理执行失败", e);
            return requests.stream().map(request ->
                    ChatCompletionResponse.builder()
                            .id("fallback")
                            .choices(List.of(ChatCompletionChoice.builder()
                                    .message(AssistantMessage.from("处理失败")).build()))
                            .build())
                    .toList();
        }
    }

    private void cacheResponses(List<ChatCompletionRequest> requests, List<ChatCompletionResponse> responses) {
        for (int i = 0; i < requests.size(); i++) {
            String cacheKey = generateCacheKey(requests.get(i));
            cacheManager.getCache("aiResponses").put(cacheKey, responses.get(i));
        }
    }

    private String generateCacheKey(ChatCompletionRequest request) {
        return DigestUtils.md5DigestAsHex(request.toString().getBytes());
    }

    private Executor getAIPoolExecutor() {
        return new ThreadPoolExecutor(
                10, 50, 60L, TimeUnit.SECONDS,
                new LinkedBlockingQueue<>(100),
                new NamedThreadFactory("AI-Worker", false),
                new ThreadPoolExecutor.CallerRunsPolicy());
    }
}
